"""
 注意：：：这个版本 测试如果不使用open ai 的大模型 返回的格式不会遵循json（应该是.with_structured_output(xxx)不支持） ，导致运行出错。所以另外基于这个改了一下提示词版本的。

reAct 大模型的一种思维模式，让大模型先思考再行动。（Reason + Act） 框架。
主要工具：
    tavily 搜索检索 ，需要配置api
    langsmit  agent 监控工具 ，需要配置api

本示例：
    使用 LangGraph 构建一个完整的“ReAct智能体工作流”。
    - Tavily 搜索工具：让模型具备网络搜索能力。
    - LangSmith（可选）：用于可视化 agent 的执行过程。

"""
import asyncio
import operator
import os
from typing import TypedDict, List, Annotated, Tuple, Union, Literal

from dotenv import load_dotenv
from langchain import hub
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent
from pydantic import Field, BaseModel

# 自动加载 .env 文件中的环境变量
load_dotenv()

# 创建 Tavily 搜索工具（可联网搜索），每次只返回一个搜索结果
tools = [TavilySearchResults(max_results=1)]

# 从 LangChain Hub 获取 ReAct 框架提示模板
# Hub 上保存了一些通用的 prompt，可以直接 pull 使用
prompt = hub.pull("wfh/react-agent-executor")
prompt.pretty_print()

# ----------------------------------------------------------
# 二、模型定义
# ----------------------------------------------------------
#  讯飞星火 "xdeepseekv32exp"
model_name = "xdeepseekv32exp"

# 定义大语言模型
llm = ChatOpenAI(api_key=os.getenv("XUNFEI_API_KEY"), openai_api_base=os.getenv("XUNFEI_API_BASE"), model=model_name, streaming=True, temperature=0.7)


# 创建 ReAct 智能体执行器
# 会自动管理“思考 → 行动 → 观察 → 再思考”的循环逻辑
agent_executor = create_react_agent(llm, tools, prompt=prompt)


# ----------------------------------------------------------
# 三、定义状态结构（State）
# ----------------------------------------------------------
# LangGraph 中每个节点执行后会更新 State（类似状态机的上下文）
# PlanExecute 用于在节点间传递任务信息
class PlanExecute(TypedDict):
    input: str  # 用户录入提示词
    plan: List[str]  # 用户录入提示词如果是大问题，会被拆分成几个小问题
    past_steps: Annotated[List[Tuple], operator.add]  # 已执行的步骤及结果
    response: str  # 最终回复给用户的答案


# ----------------------------------------------------------
# 四、定义 Pydantic 数据模型
# ----------------------------------------------------------
# 定义一个计划模型类，用于描述未来要执行的计划
class Plan(BaseModel):
    """表示一组待执行的计划步骤"""
    steps: List[str] = Field(description="需要执行的不同步骤，应该按顺序排列")


# 定义生成计划的提示模板（系统角色）
planner_prompt = ChatPromptTemplate.from_messages([
    ("system",
     """
       对于给定的目标，提出一个简单的逐步计划。
       - 每个步骤应是独立、可执行的任务；
       - 不要添加多余步骤；
       - 最后一步应能直接得到最终答案；
       - 确保每步信息足够，不要跳过。
       """),
    ("placeholder", "{messages}")  # 这里插入用户输入
])

# 构建 planner（计划生成器）
#planner = planner_prompt | ChatOpenAI(model=model_name, temperature=0).with_structured_output(Plan)
planner = planner_prompt | ChatOpenAI(model=model_name, temperature=0).bind(response_format={"type": "json_object"}) | JsonOutputParser()

# 定义最终响应模型
class Response(BaseModel):
    """表示模型的直接回答"""
    response: str


# 定义 Act 模型，表示要执行的行为
# BaseModel 是 Pydantic 的基类，继承它后，类中的每个字段都会自动生成类型校验，序列号能力
# Union 是 Python 的类型注解，表示这个字段的值可以是Union中任意类型之一。即它既可以是直接回应（Response），也可以是新的计划（Plan）
# Field(...) 是 Pydantic 提供的函数，用来给字段添加元数据（例如描述、默认值、约束条件）,Field 会作为输出约束给大模型（langchain会将它转换为提示词）
class Act(BaseModel):
    """要执行的行为"""
    # action: Union[Response, Plan] = Field(description="要执行的行为。如果要回应用户，使用Response。 如果需要进一步使用工具获取答案，使用Plan。")
    action: Union[Response, Plan] = Field(description="若已可回答则返回Response，否则返回新的Plan。")


# 定义“重新规划”的提示模板
replanner_prompt = ChatPromptTemplate.from_template(
    """对于给定的目标，提出一个简单的逐步计划。这个计划应该包含独立的任务，如果正确执行将得出正确的答案。不要添加任何多余的步骤。最后一步的结果应该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。

你的目标是：
{input}

你的原计划是：
{plan}

你目前已完成的步骤是：
{past_steps}

相应地更新你的计划。如果不需要更多步骤并且可以返回给用户，那么就这样回应。如果需要，填写计划。只添加仍然需要完成的步骤。不要返回已完成的步骤作为计划的一部分。"""
)

# 构建 replanner（重新规划器）
replanner = replanner_prompt | ChatOpenAI(model=model_name).with_structured_output(Act)
#replanner = replanner_prompt | ChatOpenAI(model=model_name).bind(response_format={"type": "json_object"}) | JsonOutputParser()



# ----------------------------------------------------------
# 五、定义工作流节点函数
# ----------------------------------------------------------
async def main():
    # ---------- 1️⃣ 计划节点 ----------
    async def plan_step(state: PlanExecute):
        """根据用户输入生成计划"""
        plan = await planner.ainvoke({"messages": [("user", state["input"])]})
        print(f"生成的计划: {plan}")
        return {"plan": plan["steps"]}

    # ---------- 2️⃣ 执行节点 ----------
    async def execute_step(state: PlanExecute):
        """执行计划中的第一个步骤"""
        plan = state["plan"]
        # Python 中非常典型的一种写法，用来把一个列表里的步骤转成带序号的多行字符串。
        # enumerate(plan) 会把plan变成一个“带索引”的可迭代对象：
        plan_str = "\n".join(f"{i + 1}.{step}" for i, step in enumerate(plan))
        task = plan[0]
        task_formatted = f"""以下是计划：
{plan_str}\n\n 你的任务是执第1步任务，{task}"""
        agent_response = await agent_executor.ainvoke({"messages": [("user", task_formatted)]})

        # 将执行结果追加到 past_steps
        return {
            # + 操作 相当于在旧列表后追加一个新的元素，但不改变原列表（创建一个新的）
            "past_steps": state["past_steps"] + [(task, agent_response["messages"][-1].content)]
        }

    # ---------- 3️⃣ 重新规划节点 ----------
    async def replan_step(state: PlanExecute):
        """根据已执行的步骤决定是否继续"""
        output = await replanner.ainvoke(state)
        # 若 replanner 返回 Response，则任务结束
        if isinstance(output.action, Response):
            return {"response": output.action.response}
        # 否则返回新的计划继续执行
        else:
            return {"plan": output.action.steps}

    # ---------- 4️⃣ 条件判断 ----------
    def should_end(state: PlanExecute) -> Literal["agent", "__end__"]:
        """判断是否结束"""
        if "response" in state and state["response"]:
            return "__end__"
        else:
            return "agent"

    # ----------------------------------------------------------
    # 六、构建 LangGraph 工作流
    # ----------------------------------------------------------
    # 定义状态图，PlanExecute 为全局共享状态结构
    workflow = StateGraph(PlanExecute)

    # 添加节点（即每个逻辑阶段）

    workflow.add_node("planner", plan_step)  # 添加计划节点
    workflow.add_node("agent", execute_step)  # 添加执行步骤节点
    workflow.add_node("replan", replan_step)  # 添加重新计划节点

    # 定义节点间的有向边（执行路径）
    workflow.add_edge(START, "planner")  # 开始 → 计划
    workflow.add_edge("planner", "agent")  # 计划 → 执行
    workflow.add_edge("agent", "replan")  # 执行 → 重新规划

    # 条件边：判断是否继续还是结束
    workflow.add_conditional_edges("replan", should_end)

    # 编译生成可运行的 LangGraph 应用
    app = workflow.compile()

    # 可视化输出状态图结构（Mermaid 格式）
    graph_png = app.get_graph().draw_mermaid_png()
    with open("llm_langgraph_1_agent工作流.png", "wb") as f:
        f.write(graph_png)

    # ----------------------------------------------------------
    # 七、运行工作流
    # ----------------------------------------------------------
    config = {"recursion_limit": 50}  # 设置递归上限防止死循环
    # 输入参数
    inputs = {"input": "2024年巴黎奥运会100米自由泳决赛冠军的家乡是哪里？请用中文答复"}

    # 异步流式运行工作流（可实时看到状态变化）
    async for event in app.astream(inputs, config=config):
        for k, v in event.items():
            if k != "__end__":
                print(v)


# 运行异步函数
asyncio.run(main())
